import hmac
import hashlib
import time
import struct
import base64
import os


# 密钥生成（符合 Google Authenticator 标准）
def generate_secret(length=16):
    if length % 8 != 0:
        raise ValueError("密钥长度必须是8的倍数")
    random_bytes = os.urandom(length)
    return base64.b32encode(random_bytes).decode('utf-8').rstrip('=')

# TOTP 计算核心
def compute_totp(secret, time_step=30, digits=6, hash_alg=hashlib.sha1):
    # 处理 Base32 填充问题
    missing_padding = (8 - (len(secret) % 8)) % 8
    secret += '=' * missing_padding
    key = base64.b32decode(secret, casefold=True)

    counter = int(time.time() // time_step)
    counter_bytes = struct.pack(">Q", counter)

    hmac_digest = hmac.new(key, counter_bytes, hash_alg).digest()
    offset = hmac_digest[-1] & 0x0F
    binary_code = struct.unpack(">I", hmac_digest[offset:offset+4])[0] & 0x7FFFFFFF

    return str(binary_code % (10 ** digits)).zfill(digits)

# 验证逻辑（允许时间偏移）
def verify_totp(secret, code, window=1):
    current_time = time.time()
    for i in range(-window, window+1):
        expected_code = compute_totp(
            secret,
            time_step=30,
            digits=len(code),
            hash_alg=hashlib.sha1
        )
        if hmac.compare_digest(expected_code, code):
            return True
    return False

# 示例用法
if __name__ == "__main__":
    # 生成新密钥（需保存并与客户端共享）
    # secret = generate_secret()
    # print(f"密钥（Base32）: {secret}")

    # 生成当前验证码
    current_code = compute_totp('qkqzu5ni6pgwavkhcmparxhihobzdvnd')
    print(f"当前验证码: {current_code}")

    # 验证测试
    # test_code = input("输入验证码: ").strip()
    # print("验证结果:", verify_totp(secret, test_code))
